#include <dicom/storage/encrypt.h>
#include <dicom/storage/decrypt.h>
#include <dicom/storage/encrypt_stream.h>
#include <gwbase/base/util.h>
#include <gwbase/base/test_harness.h>
#include <gwbase/base/test_util.h>
#include <gwbase/encrypt/compress_helper.h>
#include <surbasic/AESHelper.h>
#include <gwbase/encrypt/digest.h>

using namespace dicom;

const std::string key =
    "defc580a5dccfd05a7f2b5433fea3283a5b962627189b17ec3a6ae5239462c17";

class EncrytStreamTest {};

void CheckEncryptStreamN(size_t n, bool random, unsigned char byte) {
  Status s;
  const std::string src_file = StringPrintf("test_%ld.dat", n);
  const std::string expect_out = StringPrintf("test.expect_%ld.dat", n);
  const std::string out_file = StringPrintf("test.out_%ld.dat", n);
  Random rand;
  CreateTestFile(src_file, n, random, byte);

  HashEntry hash1;
  FileReader r;
  FileWriter w;

  s = r.Open(src_file);
  EXPECT_TRUE(s.ok());
  s = w.Open(out_file);
  EXPECT_TRUE(s.ok());
  EncryptStream es(key, &w);

  char buf[4096];
  for (;;) {
    int n = r.Read(buf, 4096);
    if (n == 0) break;
    EXPECT_TRUE(n > 0);
    int ret = es.Write(buf, n);
    EXPECT_TRUE(ret == n);
  }

  int ret = es.Finalize();
  EXPECT_TRUE(ret == 0);
  es.Close();
  r.Close();
  es.GetHash(&hash1);

  HashEntry hash2;
  s = dicom::ZipAndEncryptFile(key, src_file, expect_out, &hash2);
  EXPECT_TRUE(s.ok());

  EXPECT_TRUE(CheckSameFile(expect_out, out_file));
  EXPECT_TRUE(HashEntryEqual(hash1, hash2));

  LOG_INFO << "SIZE: " << n << " PASSED";
}

TEST(EncrytStreamTest, EncrytStreamAllZeroFileHashTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptStreamN(i, false, '\0');
  }

  CheckEncryptStreamN(1000, false, '\0');
  CheckEncryptStreamN(1024, false, '\0');
  CheckEncryptStreamN(1024 * 1024, false, '\0');
}

TEST(EncrytStreamTest, EncryptStreamRandomFileHashTest) {
  for (size_t i = 0; i < 129; ++i) {
    CheckEncryptStreamN(i, true, '\0');
  }

  CheckEncryptStreamN(1000, true, '\0');
  CheckEncryptStreamN(1024, true, '\0');
  CheckEncryptStreamN(1024 * 1024, true, '\0');
}

int main() { RunAllTests(); }
